import sys
import math
import torch
import torch.nn as nn

import utils,function

class SOAS_Local_Info(nn.Module):
    def __init__(self,input_dim,hidden_dim,latent_dim,act_dim,latent_measure="pdf",lstm_layer_num=3,cut_layer_num=3,latent_layer_num=3,act_layer_num=3,
    ahead_size=80,past_size=80):
        super(SOAS_Local_Info,self).__init__()
        
        # 保存模型参数
        self.input_dim=input_dim
        self.hidden_dim=hidden_dim
        self.latent_dim=latent_dim
        self.act_dim=act_dim
        self.latent_measure=latent_measure
        self.lstm_layer_num=lstm_layer_num
        self.cut_layer_num=cut_layer_num
        self.latent_layer_num=latent_layer_num
        self.act_layer_num=act_layer_num
        self.ahead_size=ahead_size
        self.past_size=past_size

        # 初始化LSTM层
        # state的输入到LSTM之前应该先转换维度为hidden_dim
        self.to_lstm=nn.Linear(self.input_dim,self.hidden_dim)
        # 初始化LSTM层
        self.lstm_layers=[nn.LSTMCell(self.hidden_dim,self.hidden_dim) for i in range(self.lstm_layer_num)]
        # 将所有LSTM层保存在一个nn.Sequential类对象中
        self.lstm_layers_seq=nn.Sequential(*(self.lstm_layers))
        
        # 初始化预测截断信号的网络
        self.cut_layers=[nn.Linear(self.past_size*self.hidden_dim,self.hidden_dim),nn.ReLU()]
        for i in range(self.cut_layer_num-1):
            self.cut_layers.append(nn.Linear(self.hidden_dim,self.hidden_dim))
            self.cut_layers.append(nn.ReLU())
        self.cut_layers.append(nn.Linear(self.hidden_dim,1))
        self.cut_layers.append(nn.Sigmoid())
        self.cut_net=nn.Sequential(*(self.cut_layers))

        # 初始化Latent推断网络
        self.latent_layers=[]
        for i in range(self.latent_layer_num-1):
            self.latent_layers.append(nn.Linear(self.latent_dim+self.hidden_dim,self.latent_dim+self.hidden_dim))
            self.latent_layers.append(nn.ReLU())
        if self.latent_measure=='pdf':
            self.latent_layers.append(nn.Linear(self.latent_dim+self.hidden_dim,self.latent_dim*2))
        else:
            self.latent_layers.append(nn.Linear(self.latent_dim+self.hidden_dim,self.latent_dim))
        self.latent_net=nn.Sequential(*(self.latent_layers))

        # 预测Action的网络，这里应当是输出action在各个分量上的分布，假设为高斯分布
        self.action_layers=[]
        for i in range(self.act_layer_num-1):
            self.action_layers.append(nn.Linear(self.latent_dim+self.hidden_dim,self.latent_dim+self.hidden_dim))
            self.action_layers.append(nn.ReLU())
        # 输出的是分布的期望值以及分布的对数方差
        self.action_layers.append(nn.Linear(self.latent_dim+self.hidden_dim,self.act_dim*2))
        self.action_net=nn.Sequential(*(self.action_layers))

    def forward_train(self,states,actions,slope=10):
        # 用于完整的重构一条轨迹出来
        # states是输入的状态序列
        # 获取BatchSize
        batchSize=states.shape[0]
        # 获取设备
        device=states.device
        # 用于保存所有iter_hiddens的列表
        all_iter_hiddens=[]
        # 初始化隐藏状态
        iter_hiddens=[utils.get_lstm_initial_state(batchSize,self.hidden_dim,device) for i in range(self.lstm_layer_num)]
        # 初始化top_feats
        top_feats=[torch.zeros((batchSize,self.hidden_dim)).to(device) for i in range(self.past_size)]
        # 计算截断填充特征的cut信号
        prior_cut_train=self.cut_net(torch.cat(top_feats[-self.past_size:],dim=1).detach())
        prior_cut=prior_cut_train.detach()
        # cut4train里面保存的是用于训练的cut张量，cuts保存的cut预测信号是经过detach的张量
        cuts4train=[prior_cut_train]
        cuts=[prior_cut]
        # 获取轨迹的长度
        trajLen=states.shape[1]

        # 扫描轨迹生成feats，扫描过程中，应当生成在实际cut预测器预测的截断信号作用下生成的feats
        for i in range(trajLen):
            # 计算当前时间步的隐藏状态
            now_hiddens=[]
            # 初始化当前的输入为对应的state
            state_i=self.to_lstm(states[:,i])
            # 循环遍历每一个LSTM_Cell
            for j in range(self.lstm_layer_num):
                hidden_j=self.lstm_layers[j](state_i,(iter_hiddens[j][0]*prior_cut,iter_hiddens[j][1]*prior_cut))
                now_hiddens.append(hidden_j)
                state_i=hidden_j[0]
            iter_hiddens=now_hiddens
            # 保存iter_hiddens
            all_iter_hiddens.append(iter_hiddens)
            # top_feats存放的是hidden state不包括cell state
            top_feats.append(iter_hiddens[-1][0])
            # 更新prior_cut信号以及cut信号
            prior_cut_train=self.cut_net(torch.cat(top_feats[-self.past_size:],dim=1).detach())
            prior_cut=prior_cut_train.detach()
            cuts4train.append(prior_cut_train)
            cuts.append(prior_cut)
        
        # 初始化保存Latent Variable的list对象
        latent_codings=[]
        latent_mus=[]
        latent_logvars=[]
        latents=[]
        # 扫描生成潜在变量
        prior_latent=torch.zeros((batchSize,self.latent_dim)).to(device)
        # 扫描完整的轨迹长度
        for i in range(trajLen):
            # 用于预测latent的输入向量，这个地方的逻辑还要斟酌一下
            feat4latent=torch.cat([prior_latent*cuts[i],top_feats[i+self.past_size]],dim=1)
            # 预测当前时间步的latent结果
            latent_coding=self.latent_net(feat4latent)
            # 保存latent网络输出的结果
            latent_codings.append(latent_coding)
            # 拆分得到mu和log_var
            mu,log_var=torch.split(latent_coding,self.latent_dim,dim=1)
            # 追加latent_mu和latent_logvar
            latent_mus.append(mu)
            latent_logvars.append(log_var)
            # 采样得到当前时间步的latent
            latent=function.gaussian_sample(mu,log_var)
            # 保存latent并更新以便下一个时间步的工作
            latents.append(latent)
            prior_latent=latent

        # 保存输出的动作分布的参数
        actions_mu=[]
        actions_logvar=[]
        actions_logprob=[]
        # 开始预测动作分布
        for i in range(trajLen):
            # 获取预测需要的两个张量
            latent_i=latents[i]
            feat_i=top_feats[i+self.past_size]
            # 生成用于预测的张量
            feat4action=torch.cat([latent_i,feat_i],dim=1)
            # 输出动作的预测编码
            action=self.action_net(feat4action)
            # split输出内容
            act_mu,act_logvar=torch.split(action,self.act_dim,dim=1)
            # 添加到记录结果的list中
            actions_mu.append(act_mu)
            actions_logvar.append(act_logvar)
            # 获取action
            action_i=actions[:,i]
            # 计算action的logprob
            act_logprob=torch.exp(-torch.pow(action_i-act_mu,2)/(2*torch.exp(act_logvar)))/math.sqrt(2*math.pi)/torch.exp(0.5*act_logvar)
            # 追加action的对数似然
            actions_logprob.append(act_logprob)

        # 扫描轨迹，这次生成用于训练cut_net参数的监督信息
        cut_labels=[torch.zeros((states.shape[0],1)).to(device)]

        # 接下来生成cuts4train里保存的非detach的cut信号的监督信息
        # 分为两个部分，一个部分是生成完全阻断当前时间步情况下，看ahead_size个时间步的action重构情况
        # 另一个部分是完全不阻断当前时间步的情况下，看ahead_size个时间步的action重构情况

        # 首先写完全阻断当前时间步的情况吧
        # 保存每一个时间步的向前看的action重构张量，最后一个时间步不处理，因为再往前看就没有时间步了
        absCut_act_logProb=[]
        for i in range(trajLen-1):
            # 生成i时间步向前看的top_feats列表
            i_top_feats=top_feats[i+1:i+self.past_size+1]
            # 保存i时间步向前看重构的action_logProb
            i_ahead_logProbs=[]
            # 初始化时间步i这里的cut信号为0
            prior_cut=0
            # 既然已经是0了，大可以初始化i时间步的latent为0张量
            prior_latent=torch.zeros((batchSize,self.latent_dim)).to(device)
            # 同理，讨论的是从i时间步开始完全cut，可以初始化iter_hiddens为全0的张量
            iter_hiddens=[utils.get_lstm_initial_state(batchSize,self.hidden_dim,device) for i in range(self.lstm_layer_num)]
            # 这里开始是向前看的遍历逻辑
            for j in range(i+1,min(i+self.ahead_size+1,trajLen)):
                # 用于保存j时间步的hiddens
                now_hiddens=[]
                # 转换j时间步的observation的维数
                state_j=self.to_lstm(states[:,j])
                # 遍历各个LSTM层
                for k in range(self.lstm_layer_num):
                    hidden_k=self.lstm_layers[k](state_j,(iter_hiddens[k][0]*prior_cut,iter_hiddens[k][1]*prior_cut))
                    now_hiddens.append(hidden_k)
                    state_j=hidden_k[0]
                # 更新iter_hiddens
                iter_hiddens=now_hiddens
                # 更新top_feat
                top_feat=iter_hiddens[-1][0]
                # j时间步top_feat添加到i_top_feats
                i_top_feats.append(top_feat)
                # 拼接出计算j时间步latent的特征
                feat4latent=torch.cat([prior_latent*prior_cut,top_feat],dim=1)
                # 计算j时间步的latent_coding
                latent_coding=self.latent_net(feat4latent)
                # 计算实际的latent
                mu,log_var=torch.split(latent_coding,self.latent_dim,dim=1)
                now_latent=function.gaussian_sample(mu,log_var)
                # 更新prior_cut信号
                prior_cut=self.cut_net(torch.cat(i_top_feats[-self.past_size:],dim=1).detach())
                # 拼接生成计算action的特征向量
                feat4action=torch.cat([now_latent,top_feat],dim=1)
                action=self.action_net(feat4action)
                # split输出内容
                act_mu,act_logvar=torch.split(action,self.act_dim,dim=1)
                # 计算j时间步的logProb
                action_j=actions[:,j]
                # 计算action的logProb
                act_logProb=torch.exp(-torch.pow(action_j-act_mu,2)/(2*torch.exp(act_logvar)))/math.sqrt(2*math.pi)/torch.exp(0.5*act_logvar)
                # 将这个logProb添加到i_ahead_logProbs
                i_ahead_logProbs.append(act_logProb)
            # 将从i时间步向前看的所有logProbs保存到absCut_act_logProb
            absCut_act_logProb.append(i_ahead_logProbs)
        
        # 再看完全不进行cut的情况
        # 保存每一个时间步的向前看的action重构张量，最后一个时间步不处理，因为再往前看就没有时间步了
        noCut_act_logProb=[]
        for i in range(trajLen-1):
            # 生成i时间步向前看的top_feats列表
            i_top_feats=top_feats[i:i+self.past_size+1]
            # 保存i时间步向前看重构的action_logProb
            i_ahead_logProbs=[]
            # 初始化时间步i这里的cut信号为1
            prior_cut=1
            # 取出i时间步的latent张量
            prior_latent=latents[i]
            # 获取i时间步生成的各层LSTM生成的hiddens
            iter_hiddens=all_iter_hiddens[i]
            # 下面开始向前看的遍历逻辑
            for j in range(i+1,min(i+self.ahead_size+1,trajLen)):
                # 用于保存j时间步的hiddens
                now_hiddens=[]
                # 转换j时间步的observation的维数
                state_j=self.to_lstm(states[:,j]) 
                # 遍历各个LSTM层
                for k in range(self.lstm_layer_num):
                    hidden_k=self.lstm_layers[k](state_j,(iter_hiddens[k][0]*prior_cut,iter_hiddens[k][1]*prior_cut))
                    now_hiddens.append(hidden_k)
                    state_j=hidden_k[0]
                # 更新iter_hiddens
                iter_hiddens=now_hiddens
                # 更新top_feat
                top_feat=iter_hiddens[-1][0]
                # j时间步top_feat添加到i_top_feats
                i_top_feats.append(top_feat)
                # 拼接出计算j时间步latent的特征
                feat4latent=torch.cat([prior_latent*prior_cut,top_feat],dim=1)
                # 计算j时间步的latent_coding
                latent_coding=self.latent_net(feat4latent)
                # 计算实际的latent
                mu,log_var=torch.split(latent_coding,self.latent_dim,dim=1)
                now_latent=function.gaussian_sample(mu,log_var)
                # 更新prior_cut信号
                prior_cut=self.cut_net(torch.cat(i_top_feats[-self.past_size:],dim=1).detach())
                # 拼接生成计算action的特征向量
                feat4action=torch.cat([now_latent,top_feat],dim=1)
                action=self.action_net(feat4action)
                # split输出内容
                act_mu,act_logvar=torch.split(action,self.act_dim,dim=1)
                # 计算j时间步的logProb
                action_j=actions[:,j]
                # 计算action的logProb
                act_logProb=torch.exp(-torch.pow(action_j-act_mu,2)/(2*torch.exp(act_logvar)))/math.sqrt(2*math.pi)/torch.exp(0.5*act_logvar)
                # 将这个logProb添加到i_ahead_logProbs
                i_ahead_logProbs.append(act_logProb)
            # 将从i时间步向前看的所有logProbs保存到noCut_act_logProb
            noCut_act_logProb.append(i_ahead_logProbs)

        # noCut_act_logProb, absCut_act_logProb都是list类型，其中的元素也是list类型，元素中的元素是一个时间步的action重构
        # 这两个list的长度都是trajLen-1，需要构造的cut_labels列表长度应该是trajLen+1，已经有了1个label，最末尾还应该补上一个0的label
        # 接下来应该给cut_labels增加trajLen-1个元素，并且在最末尾追加0 label
        for i in range(len(absCut_act_logProb)):
            # 针对i时间步，对向前看时间窗口内的action的对数似然求平均，去除向前看的步数的影响
            absCut_mean_logProb=sum(absCut_act_logProb[i])/len(absCut_act_logProb[i])
            noCut_mean_logProb=sum(noCut_act_logProb[i])/len(noCut_act_logProb[i])
            # 计算不进行cut和完全cut的logProb之间的差值
            res_logProb=noCut_mean_logProb-absCut_mean_logProb
            # 对action各个分量的logProb进行求平均
            mean_res_logProb=torch.mean(res_logProb,dim=1,keepdim=True)
            # 对整个张量进行先放大后Sigmoid的操作生成合适的cut_net的监督标签
            i_cut_label=torch.sigmoid(mean_res_logProb*slope).detach()
            cut_labels.append(i_cut_label)
        cut_labels.append(torch.zeros(batchSize,1).to(device))

        return actions_logprob,latent_mus,latent_logvars,cuts4train,cut_labels






